import sys
import requests
from bs4 import BeautifulSoup
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QLineEdit, QPushButton, QVBoxLayout, QMessageBox


class SpiderUI(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        # 创建控件
        self.keyword_label = QLabel('关键字:')
        self.keyword_input = QLineEdit()
        self.count_label = QLabel('爬取数量:')
        self.count_input = QLineEdit()
        self.crawl_button = QPushButton('开始爬取')
        self.export_button = QPushButton('导出数据')
        self.crawl_button.clicked.connect(self.crawl_data)
        self.export_button.clicked.connect(self.export_data)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.keyword_label)
        layout.addWidget(self.keyword_input)
        layout.addWidget(self.count_label)
        layout.addWidget(self.count_input)
        layout.addWidget(self.crawl_button)
        layout.addWidget(self.export_button)

        self.setLayout(layout)
        self.setWindowTitle('爬虫界面')
        self.setGeometry(300, 300, 300, 200)
        self.show()

    def crawl_data(self):
        keyword = self.keyword_input.text()
        try:
            count = int(self.count_input.text())
        except ValueError:
            QMessageBox.warning(self, '错误', '请输入有效的爬取数量')
            return
        if not keyword:
            QMessageBox.warning(self, '错误', '请输入关键字')
            return
        url = f'https://www.baidu.com/s?wd={keyword}'
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'}
        try:
            response = requests.get(url, headers=headers)
            response.raise_for_status()
            soup = BeautifulSoup(response.text, 'html.parser')
            results = soup.find_all('h3')
            self.data = [result.get_text() for result in results[:count]]
            QMessageBox.information(self, '提示', f'成功爬取 {len(self.data)} 条数据')
        except requests.RequestException as e:
            QMessageBox.warning(self, '错误', f'请求出错: {e}')
        except Exception as e:
            QMessageBox.warning(self, '错误', f'发生未知错误: {e}')

    def export_data(self):
        if hasattr(self, 'data'):
            try:
                with open('data.txt', 'w', encoding='utf-8') as f:
                    for item in self.data:
                        f.write(item + '\n')
                QMessageBox.information(self, '提示', '数据导出成功')
            except Exception as e:
                QMessageBox.warning(self, '错误', f'导出数据出错: {e}')
        else:
            QMessageBox.warning(self, '错误', '请先进行爬取')


if __name__ == '__main__':
    app = QApplication(sys.argv)
    spider_ui = SpiderUI()
    sys.exit(app.exec_())
